import torch
from cnn_create import CNN
from dataset_create import TraficSignDataSet
from torchvision import transforms
import torch.nn as nn

# Params
num_epochs = 5
num_class = 43
batch_size = 50
learning_rate = 0.001

# Load the LeNet model
net = CNN(num_class)
print(net)
net.load_state_dict(torch.load('net.ckpt'))
net.eval()
print('Net model load success.')

# Load the test dataset
test_dataset = TraficSignDataSet(need_train=False, transform=transforms.ToTensor())
test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                          batch_size=batch_size,
                                          shuffle=True)
print('Test dataset load success.')

# Test
correct = 0
total = 0
print('Testing ...')
with torch.no_grad():
    for data in test_loader:
        image, label = data['image'], data['label']
        outputs = net(image)
        _, predicted = torch.max(outputs.data, 1)
        total += label.size(0)
        correct += (predicted == label).sum().item()

print('Accuracy of all the test dataset : {:.2f} %'.format(100 * correct / total))
